import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from adjustText import adjust_text
from matplotlib.colors import LinearSegmentedColormap
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import warnings
warnings.filterwarnings("ignore", message=".*FancyArrowPatch.*")

# 设置中文字体（例如黑体 SimHei）
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或者 ['Microsoft YaHei']
# 解决保存图像时负号'-'显示为方块的问题
plt.rcParams['axes.unicode_minus'] = False
# 创建示例数据（假设有5个不同模型）
models_data = {
    'Model': ['Zhu', 'Huang', 'Le K H', 'Ghaish', '本章'],
    'F1_Score': [93.39, 94.01, 95.02, 95.69, 95.53],  # F1分数
    'Parameters(M)': [16.59, 1.13, 4.93, 0.66, 0.18],  # 参数量（百万）
    'FLOPs(B)': [224.09, 77.64, 111.5, 87.7, 34.4]
}

df = pd.DataFrame(models_data)

# 创建图形和坐标轴
plt.figure(figsize=(10, 6))

# 创建自定义颜色映射（从蓝色到红色）
colors = plt.cm.viridis(np.linspace(0, 1, len(df)))

# 绘制散点图
scatter = plt.scatter(
    df['Parameters(M)'],
    # df['FLOPs(B)'],
    df['F1_Score'],
    c=colors,
    s=100,  # 点的大小
    alpha=0.7,  # 透明度
    edgecolors='black',  # 边缘颜色
    linewidth=1)  # 边缘线宽

# 添加模型名称标签
# for i, row in df.iterrows():
#     text = plt.annotate(row['Model'],
#                         # (row['Parameters(M)'], row['F1_Score']),
#                         (row['FLOPs(B)'], row['F1_Score']),
#                         xytext=(5, 5),  # 标签偏移量
#                         textcoords='offset points',
#                         fontsize=12,
#                         alpha=0.8)

# 设置坐标轴标签和标题
plt.xlabel('模型参数量 (百万)', fontsize=12, fontweight='bold')
# plt.xlabel('浮点运算数', fontsize=12, fontweight='bold')
plt.ylabel('MF1分数', fontsize=12, fontweight='bold')
plt.title('MF1分数与参数量对比分析', fontsize=14, fontweight='bold')
# plt.title('MF1分数与浮点运算数对比分析', fontsize=14, fontweight='bold')

# 添加网格线
plt.grid(True, alpha=0.3, linestyle='--')

# 设置坐标轴范围（留出一些边距）
plt.xlim(0, df['Parameters(M)'].max() * 1.1)
# plt.xlim(0, df['FLOPs(B)'].max() * 1.1)
plt.ylim(df['F1_Score'].min() * 0.95, df['F1_Score'].max() * 1.05)

# 添加颜色条
# cbar = plt.colorbar(scatter)
# cbar.set_label('模型复杂度', rotation=270, labelpad=15)

# 添加趋势线（可选）
z = np.polyfit(df['Parameters(M)'], df['F1_Score'], 1)
# z = np.polyfit(df['FLOPs(B)'], df['F1_Score'], 1)
p = np.poly1d(z)
plt.plot(df['Parameters(M)'], p(df['Parameters(M)']), "r--", alpha=0.5, label='趋势线')
# plt.plot(df['FLOPs(B)'], p(df['FLOPs(B)']), "r--", alpha=0.5, label='趋势线')

# 添加图例
plt.legend()

# 自动调整布局
plt.tight_layout()

# 1. 创建一个空列表来保存所有的标签对象
texts = []
for i, model in enumerate(df['Model']):
    # 2. 添加标签，但不立即显示，将返回的Text对象存入列表
    text = plt.text(df['Parameters(M)'].iloc[i], df['F1_Score'].iloc[i], model, fontsize=12)
    # text = plt.text(df['FLOPs(B)'].iloc[i], df['F1_Score'].iloc[i], model, fontsize=12)
    texts.append(text)

# 3. 魔法发生在这里：自动调整所有标签的位置以避免重叠
adjust_text(texts,
            arrowprops=dict(arrowstyle='->', color='blue', lw=3, alpha=0.6, shrinkA=15),
            force_text=0.9,
            force_points=0.9,  # 很强的排斥力
            expand_points=(2.0, 2.0),
            expand_text=(2.0, 2.0),
            avoid_text=True,  # 强制避免文本重叠
            avoid_points=True,  # 强制避免与点重叠
            save_steps=False,  # 调试时可设为True查看过程
            lim=2000
            )
# 显示图形
plt.show()

# 可选：保存图片
plt.savefig('图15MF1分数与参数量对比分析图.jpg', dpi=300, bbox_inches='tight')
